/****************************************************************************
*
*   StickyIME.cpp
*   DLL Loader
*
*   By Cliff Spradlin (2/12/2012)
*
***/

#include <Windows.h>
#include <assert.h>
#include <Tlhelp32.h>
#include <Shlwapi.h>


/****************************************************************************
*
*   Private functions
*
***/

//===========================================================================
static void Shutdown () {
    DWORD currentProcess = GetCurrentProcessId();
    HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS | TH32CS_SNAPTHREAD, 0);

    PROCESSENTRY32W process;
    process.dwSize = sizeof(process);
    BOOL success = Process32First(snapshot, &process);
    
    while (success) {
        if (process.th32ProcessID != currentProcess && wcscmp(process.szExeFile, L"StickyIME.exe") == 0) {
            THREADENTRY32 thread;
            thread.dwSize = sizeof(thread);
            success = Thread32First(snapshot, &thread);

            while (success) {
                if (thread.th32OwnerProcessID == process.th32ProcessID) {
                    PostThreadMessage(thread.th32ThreadID, WM_QUIT, 0, 0);
                }
                success = Thread32Next(snapshot, &thread);
            }

            HANDLE processHandle = OpenProcess(SYNCHRONIZE, false, process.th32ProcessID);
            if (processHandle != NULL) {
                WaitForSingleObject(processHandle, INFINITE);
                CloseHandle(processHandle);
            }
        }
        success = Process32Next(snapshot, &process);
    }

    CloseHandle(snapshot);
}


/****************************************************************************
*
*   Exports
*
***/

//===========================================================================
int main (int argc, const char * argv[]) {
    // first, change current working directory to app directory
    // StickyIME is run automatically at startup via a registry entry, but
    // unfortunately the working directory is not the app directory.
    wchar_t path[MAX_PATH];
    GetModuleFileName(NULL, path, MAX_PATH);
    PathRemoveFileSpec(path);
    SetCurrentDirectory(path);

    // cleanup previous runs
    SendMessageTimeout(HWND_BROADCAST, WM_NULL, 0, 0, SMTO_ABORTIFHUNG|SMTO_NOTIMEOUTIFNOTHUNG, 1000, NULL);

    if (argc == 2 && strcmp(argv[1], "shutdown") == 0) {
        // called by uninstaller
        Shutdown();
        return 0;
    }

#ifndef _M_X64
    {
        // start 64-bit version
        // we need both versions running so that we can inject 
        // into both 32-bit and 64-bit processes
        STARTUPINFO info;
        PROCESS_INFORMATION pi; 

        ZeroMemory(&info, sizeof(info));
        info.cb = sizeof(info);
        CreateProcess(L"x64\\StickyIME.exe", NULL, NULL, NULL, FALSE, 0, NULL, NULL, &info, &pi);
        CloseHandle(pi.hProcess);
        CloseHandle(pi.hThread);
    }
#endif 

    HMODULE dll = LoadLibrary(L"StickyIME.dll");
    HOOKPROC callWndProc = (HOOKPROC) GetProcAddress(dll, "CallWndProc");

    // set global hooks
    HHOOK callWndHook = SetWindowsHookEx(WH_CALLWNDPROC, callWndProc, dll, 0);

    // wait until shutdown
    MSG msg;
    while (GetMessage(&msg, NULL, 0, 0) > 0) {
        TranslateMessage(&msg);
        DispatchMessage(&msg);
    }

    // unhook
    UnhookWindowsHookEx(callWndHook);
    FreeLibrary(dll);

    // free DLL in other processes
    SendMessageTimeout(HWND_BROADCAST, WM_NULL, 0, 0, SMTO_ABORTIFHUNG|SMTO_NOTIMEOUTIFNOTHUNG, 1000, NULL);
    
    return 0;
}
